#!/usr/bin/env python3
import rospy
from std_msgs.msg import Float32,Float32MultiArray
import threading
import time
import math
from pynput import keyboard
import numpy as np

front_hub=rospy.Publisher('/front_velocity',Float32,queue_size=10)
rear_pub=rospy.Publisher('/rear_velocity',Float32,queue_size=10)
joint_pubs=[
    rospy.Publisher(f'/joint_{i+1}_angle',Float32,queue_size=10) for i in range(4)
]

#Snake Info
Snake_Height=0.14
Snake_Length=0.43 #per one
#Stair Info
Stair_Height=0.18
Stair_Depth=0.22
Stair_Degree=math.degrees(math.atan(Stair_Height/Stair_Depth))

Stair_Bound=Stair_Degree/8#前履带IMU变化阈值

Third_Bound=Stair_Degree/2#关节三变化阈值

Stiar_Num=6

IMU_data=[] #用于存储 /robot/tilt_angles 的数据

joint1_to_front=0.055
joint1_to_back=0.286

front_move=0.0
back_move=0.0
Two_Roll = 0.0
Four_Roll = 0.0
One_Roll = 0.0
Three_Roll = 0.0
DELEY_TIME = 0.1  # 延迟时间，单位秒

front_motor_running=False
back_motor_running=False

move_velocity=3

joint_angles=[0.0,0.0,0.0,0.0]
is_rotating=[False,False,False,False]

keys_to_joints={
    'r':(0,1),
    'z':(0,-1),
    'a':(1,1),
    'd':(1,-1),
    't':(2,1),
    'x':(2,-1),
    'q':(3,1),
    'e':(3,-1),
}

def tilt_angles_callback(msg):
    global IMU_data
    IMU_data=msg.data
    #IMU_data[1] 保存的是前履带的倾斜角

def publish_joint_angles(pubs):
    """
    持续发布所有关节角度的线程函数
    """
    global joint_angles,is_rotating
    rate=rospy.Rate(10)
    while not rospy.is_shutdown():
        for i in range(4):
            if is_rotating[i]:
                pubs[i].publish(joint_angles[i])
        rate.sleep()

lock=threading.Lock()

def publish_velocity_time(pub,velocity,duration):
    """
    以速度执行一定时间。
    """
    rate=rospy.Rate(10)
    start_time=time.time()
    while time.time()-start_time<duration:
        pub.publish(velocity)
        rate.sleep()
    pub.publish(0.0)

def publish_velocity(pub,velocity):
    rate=rospy.Rate(10)
    pub.publish(velocity)
    rate.sleep()


def frontmotor_thread():
    global front_move, back_move,front_motor_running,back_motor_running

    while not rospy.is_shutdown():
        with lock:
            if front_move==1.0:
                if not front_motor_running:
                    publish_velocity(front_hub,move_velocity)
                    front_motor_running=True

            elif front_move==-1.0:
                if not front_motor_running:
                    publish_velocity(front_hub,-move_velocity)
                    front_motor_running=True
            else:
                if front_motor_running:
                    publish_velocity(front_hub,0)
                    front_motor_running=False

def backmotor_thread():
    global back_move,back_motor_running
    while not rospy.is_shutdown():
        with lock:
            if back_move==1.0:
                if not back_motor_running:
                    publish_velocity(rear_pub,move_velocity)
                    back_motor_running=True

            elif back_move==-1.0:
                if not back_motor_running:
                    publish_velocity(rear_pub,-move_velocity)
                    back_motor_running=True
            else:
                if back_motor_running:
                    publish_velocity(rear_pub,0)
                    back_motor_running=False

def moving_average_filter(data_list,window_size):
    if len(data_list)<window_size:
        return np.mean(data_list)
    return np.mean(data_list[-window_size:])

def low_pass_filter(new_value,last_value,alpha):
    return alpha*new_value+(1-alpha)*last_value

def ClimbStairs():
    #step1 抬头至楼梯倾斜度
    global joint_angles,joint_pubs,IMU_data,front_move, back_move

    rospy.loginfo("爬楼梯函数已启动")  # 添加日志以确认函数已启动
    while joint_angles[0]*100 > -Stair_Degree:
        joint_angles[0] -= 0.01
        joint_pubs[0].publish(joint_angles[0])
        rospy.loginfo(f"当前角度: {joint_angles[0]}")
        time.sleep(0.1)  # 可以根据实际情况调节延时
    rospy.loginfo("抬头完成")

    #step2 前进，直至前履带的IMU角度增加10度(说明前履带和楼梯碰到了)
    #count=0
    IMU_data1=IMU_data[1]
    while IMU_data1-Stair_Bound*4<IMU_data[1]:
        back_move=1
        front_move=1

    #step3 低头至摆直
    while joint_angles[0]*100<0:
        joint_angles[0]+=0.02
        joint_pubs[0].publish(joint_angles[0])
        rospy.loginfo(f"当前角度: {joint_angles[0]}")
        time.sleep(0.1)
    rospy.loginfo("拉直完成")
    
    while IMU_data[1]>-Stair_Degree:
        rospy.loginfo("爬行中")

    rospy.loginfo("已完全平行与楼梯，记录当前IMU值")
    IMU_data1=IMU_data[1]

    #step4 直行至前履带IMU变化（说明前履带末端已越过顶层阶梯）
    count=0
    while count<=5:
        if IMU_data1+Stair_Bound<IMU_data[1]:
            count=count+1
            print(f"IMU_data2+5:{IMU_data1}"+f"IMU:{IMU_data[1]}"+f"count:{count}")
        else:
            count=0
        time.sleep(0.1)
    front_move=0.0
    back_move=0.0

    #step5 开始低头至IMU>-5（只低头，不然会侧翻）
    rospy.loginfo("关节一低头")
    while IMU_data[1]<-Stair_Bound:
        joint_angles[0]+=0.02
        joint_pubs[0].publish(joint_angles[0])
        rospy.loginfo(f"当前角度: {joint_angles[0]}")
        time.sleep(0.1)
    
    #step6 前进直至IMU>5
    rospy.loginfo("前进至前履带低头")
    while IMU_data[1]<Stair_Bound:
        front_move=1.0
        back_move=1.0

    #step7 关节三低头（不然会磕到后面的电机）
    rospy.loginfo("关节三低头")
    while joint_angles[2]*100<Third_Bound:
        joint_angles[2]+=0.04
        joint_pubs[2].publish(joint_angles[2])
        time.sleep(0.1)
    
    
    time.sleep(3)
    #需要保持当前状态前进一段距离，不然上不去


    #step8 最后一步  拉直！
    rospy.loginfo("最后一步  拉直！")
    while joint_angles[0]*100>0:
        rospy.loginfo(f"当前关节一角度: {joint_angles[0]}")
        joint_angles[0]-=0.04
        joint_pubs[0].publish(joint_angles[0])
        time.sleep(0.1)

    while joint_angles[2]*100>0:
        rospy.loginfo(f"当前关节三角度: {joint_angles[2]}")
        joint_angles[2]-=0.04
        joint_pubs[2].publish(joint_angles[2])
        time.sleep(0.1)

    front_move=0.0
    back_move=0.0

def on_press(key):
    global front_move, back_move, Two_Roll, Four_Roll, One_Roll,  Three_Roll
    global joint_angles,is_rotating
    try:
        with lock:
            if key.char == 'w':
                front_move = 1.0
            elif key.char=='m':
                front_move=-1.0
            elif key.char == 's':
                back_move = 1.0
            elif key.char=='n':
                back_move=-1.0
            elif key.char=='v':
                climb_thread=threading.Thread(target=ClimbStairs)
                climb_thread.start()
            elif key.char in keys_to_joints:
                joint_idx,direction=keys_to_joints[key.char]
                joint_angles[joint_idx]+=direction*0.01
                is_rotating[joint_idx]=True
    except AttributeError:
        pass


def on_release(key):
    global front_move, back_move, back_motor_running,front_motor_running,is_rotating
    try:
        with lock:
            if key.char=='w' or key.char=='m':
                front_move=0.0
                publish_velocity(front_hub,0)
                front_motor_running=False
            elif key.char=='s' or key.char=='n':
                back_move=0.0
                publish_velocity(rear_pub,0)
                back_motor_running=False
            elif key.char in keys_to_joints:
                joint_idx,_=keys_to_joints[key.char]
                is_rotating[joint_idx]=False

    except AttributeError:
        pass


def main():
    rospy.init_node('AutoClimb',anonymous=True)

    rospy.Subscriber('/robot/tilt_angles',Float32MultiArray,tilt_angles_callback)

    front_thread=threading.Thread(target=frontmotor_thread)
    front_thread.start()

    back_thread=threading.Thread(target=backmotor_thread)
    back_thread.start()

    pub_thread=threading.Thread(target=publish_joint_angles,args=(joint_pubs,))
    pub_thread.start()
    
    listener=keyboard.Listener(on_press=on_press,on_release=on_release)
    listener.start()

    rospy.spin()

if __name__=='__main__':
    main()